import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import json
import logging
from datetime import datetime
import time

# 设置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("DroneSimulator")


# 无人机仿真类
class Quadrotor:
    def __init__(self):
        self.x = 0.0
        self.y = 0.0
        self.z = 0.0  # 高度，0表示地面
        self.roll = 0.0
        self.pitch = 0.0
        self.yaw = 0.0
        self.size = 1.0
        self.armed = False  # 无人机是否上电
        self.flying = False  # 无人机是否起飞
        self.trajectory = []  # 存储飞行轨迹
        self.log = []  # 存储飞行日志
        self.fig = plt.figure(figsize=(14, 10))
        self.ax = self.fig.add_subplot(111, projection='3d')
        self.ani = None
        self.target_points = []  # 存储目标点
        self.setup_plot()
        self.log_event("SYSTEM", "Drone simulation initialized")

    def setup_plot(self):
        """设置可视化图"""
        self.ax.set_xlim([-5, 5])
        self.ax.set_ylim([-5, 5])
        self.ax.set_zlim([0, 10])
        self.ax.set_xlabel('X (m)')
        self.ax.set_ylabel('Y (m)')
        self.ax.set_zlabel('Z (m)')
        self.ax.set_title('Quadrotor UAV MPC Control Simulation')
        self.ax.grid(True)

    def transformation_matrix(self):
        """计算从机体坐标系到世界坐标系的变换矩阵"""
        x, y, z, roll, pitch, yaw = self.x, self.y, self.z, self.roll, self.pitch, self.yaw
        cy, sy = np.cos(yaw), np.sin(yaw)
        cp, sp = np.cos(pitch), np.sin(pitch)
        cr, sr = np.cos(roll), np.sin(roll)

        dcm = np.array([
            [cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
            [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
            [-sp, cp * sr, cp * cr]
        ])

        return dcm

    def get_corners(self):
        """获取无人机四个旋翼在世界坐标系中的位置"""
        T = self.transformation_matrix()
        half_size = self.size / 2
        corners_body = np.array([
            [half_size, 0, 0],  # 右前
            [-half_size, 0, 0],  # 左后
            [0, half_size, 0],  # 前
            [0, -half_size, 0]  # 后
        ]).T

        corners_world = np.dot(T, corners_body)
        corners_world[0, :] += self.x
        corners_world[1, :] += self.y
        corners_world[2, :] += self.z

        return corners_world

    def update_plot(self, frame):
        """更新动画帧"""
        self.ax.cla()
        self.setup_plot()

        # 绘制目标点
        if self.target_points:
            targets = np.array(self.target_points)
            self.ax.scatter(targets[:, 0], targets[:, 1], targets[:, 2], c='green', marker='*', s=200,
                            label='Target Points')

            # 标记目标点编号
            for i, (x, y, z) in enumerate(targets):
                self.ax.text(x, y, z, f'T{i + 1}', color='green', fontsize=12)

        if self.flying or self.armed:
            corners = self.get_corners()
            trajectory = np.array(self.trajectory) if self.trajectory else np.array([])

            # 绘制无人机主体
            self.ax.scatter(corners[0, :], corners[1, :], corners[2, :], c='red', marker='o', s=100)
            self.ax.plot([corners[0, 0], corners[0, 1]], [corners[1, 0], corners[1, 1]], [corners[2, 0], corners[2, 1]],
                         'b-', linewidth=2)
            self.ax.plot([corners[0, 2], corners[0, 3]], [corners[1, 2], corners[1, 3]], [corners[2, 2], corners[2, 3]],
                         'g-', linewidth=2)

            # 绘制轨迹
            if len(trajectory) > 0:
                self.ax.plot(trajectory[:, 0], trajectory[:, 1], trajectory[:, 2], 'k--', linewidth=1.5,
                             label='Trajectory')

            # 标记当前位置
            self.ax.text(self.x, self.y, self.z, f'Current\n({self.x:.1f}, {self.y:.1f}, {self.z:.1f})',
                         color='blue', fontsize=10, ha='center')

            # 标记最终位置（如果已降落）
            if not self.flying and self.armed:
                self.ax.scatter(self.x, self.y, self.z, c='magenta', marker='X', s=200, label='Final Position')
                self.ax.text(self.x, self.y, self.z, f'Final\n({self.x:.1f}, {self.y:.1f}, {self.z:.1f})',
                             color='magenta', fontsize=10, ha='center')

        self.ax.legend()
        return self.ax,

    def start_animation(self):
        """启动动画"""
        self.ani = FuncAnimation(self.fig, self.update_plot, interval=100, blit=True, cache_frame_data=False)
        plt.ion()  # 开启交互模式
        plt.show(block=False)  # 非阻塞显示

    def update_state(self, x=None, y=None, z=None, roll=None, pitch=None, yaw=None):
        """更新无人机状态并记录"""
        if x is not None: self.x = x
        if y is not None: self.y = y
        if z is not None: self.z = z
        if roll is not None: self.roll = roll
        if pitch is not None: self.pitch = pitch
        if yaw is not None: self.yaw = yaw

        # 记录轨迹
        self.trajectory.append((self.x, self.y, self.z))

    def log_event(self, event_type, message, details=None):
        """记录日志事件"""
        log_entry = {
            "timestamp": datetime.now().isoformat(),
            "type": event_type,
            "message": message,
            "state": {
                "x": self.x, "y": self.y, "z": self.z,
                "roll": self.roll, "pitch": self.pitch, "yaw": self.yaw,
                "armed": self.armed, "flying": self.flying
            }
        }
        if details:
            log_entry["details"] = details

        self.log.append(log_entry)
        logger.info(f"{event_type}: {message}")

    def get_logs(self, log_type=None):
        """获取日志，可选择按类型过滤"""
        if log_type:
            return [log for log in self.log if log['type'] == log_type]
        return self.log

    def save_logs_to_file(self, filename="drone_log.json"):
        """将日志保存到JSON文件"""
        with open(filename, 'w') as f:
            json.dump(self.log, f, indent=2)
        self.log_event("SYSTEM", f"Logs saved to {filename}", {"filename": filename})
        return f"Logs saved to {filename}"


# 创建无人机实例
drone = Quadrotor()


# 控制命令函数
def arm_drone():
    """解锁无人机电机"""
    if drone.armed:
        return "Drone is already armed."

    drone.armed = True
    drone.log_event("STATE_CHANGE", "Drone armed")
    return "Drone armed."


def disarm_drone():
    """锁定无人机电机"""
    if not drone.armed:
        return "Drone is not armed."
    if drone.flying:
        return "Cannot disarm while flying."

    drone.armed = False
    drone.log_event("STATE_CHANGE", "Drone disarmed")
    return "Drone disarmed."


def takeoff():
    """无人机起飞"""
    if not drone.armed:
        return "Drone must be armed before takeoff."
    if drone.flying:
        return "Drone is already flying."

    drone.flying = True
    # 模拟起飞，飞到2米高度
    drone.update_state(z=2.0)
    drone.log_event("CONTROL", "Takeoff command", {"altitude": 2.0})
    return "Taking off to 2m altitude."


def land():
    """无人机降落"""
    if not drone.flying:
        return "Drone is not flying."

    # 模拟降落，高度降为0
    drone.update_state(z=0.0)
    drone.flying = False
    drone.log_event("CONTROL", "Land command")
    return "Landing."


def move(x=0, y=0, z=0, yaw=0):
    """移动无人机。
    x: 前后移动（米，正数向前）
    y: 左右移动（米，正数向左）
    z: 垂直移动（米，正数向上）
    yaw: 偏航角（弧度，正数顺时针）
    """
    if not drone.flying:
        return "Drone is not flying. Takeoff first."

    # 更新状态
    new_x = drone.x + x
    new_y = drone.y + y
    new_z = drone.z + z
    new_yaw = drone.yaw + yaw

    drone.update_state(x=new_x, y=new_y, z=new_z, yaw=new_yaw)
    drone.log_event("CONTROL", "Move command", {"x": x, "y": y, "z": z, "yaw": yaw})
    return f"Moving to position: x={new_x:.2f}, y={new_y:.2f}, z={new_z:.2f}, yaw={new_yaw:.2f}"


def return_to_home():
    """返回起飞点"""
    if not drone.flying:
        return "Drone is not flying."

    drone.update_state(x=0, y=0, yaw=0)
    drone.log_event("CONTROL", "Return to home command")
    return "Returning to home position."


# MPC控制器模拟
class MPCController:
    def __init__(self, drone):
        self.drone = drone
        self.target_positions = [
            (0, 0, 2),  # 起飞点
            (2, 0, 3),  # 第一个目标点
            (2, 2, 4),  # 第二个目标点
            (-2, 2, 3),  # 第三个目标点
            (-2, -2, 2),  # 第四个目标点
            (0, 0, 2)  # 返回起点
        ]
        self.drone.target_points = self.target_positions  # 设置目标点用于可视化
        self.current_target_idx = 0
        self.reached_target = False
        self.control_history = []  # 存储控制历史

    def calculate_control(self):
        """计算控制输入（简化版的MPC）"""
        if not self.drone.flying:
            return None

        target_x, target_y, target_z = self.target_positions[self.current_target_idx]

        # 当前位置
        current_x, current_y, current_z = self.drone.x, self.drone.y, self.drone.z

        # 计算误差
        error_x = target_x - current_x
        error_y = target_y - current_y
        error_z = target_z - current_z

        # 检查是否到达目标点（容差范围内）
        distance = np.sqrt(error_x ** 2 + error_y ** 2 + error_z ** 2)
        if distance < 0.2:  # 到达目标点的容差
            self.reached_target = True
            return None

        # 简化控制：按比例向目标移动
        # 在实际MPC中，这里会有更复杂的优化计算
        move_x = 0.2 * error_x / max(distance, 0.1)
        move_y = 0.2 * error_y / max(distance, 0.1)
        move_z = 0.2 * error_z / max(distance, 0.1)

        # 记录控制计算信息
        control_info = {
            "time": time.time(),
            "target_idx": self.current_target_idx,
            "target_position": (target_x, target_y, target_z),
            "current_position": (current_x, current_y, current_z),
            "errors": (error_x, error_y, error_z),
            "distance": distance,
            "control_input": (move_x, move_y, move_z, 0)
        }
        self.control_history.append(control_info)

        # 打印MPC计算信息
        print("\n" + "=" * 80)
        print("MPC CONTROL CALCULATION")
        print("=" * 80)
        print(f"Target {self.current_target_idx + 1}: {target_x:.2f}, {target_y:.2f}, {target_z:.2f}")
        print(f"Current: {current_x:.2f}, {current_y:.2f}, {current_z:.2f}")
        print(f"Errors: X={error_x:.2f}, Y={error_y:.2f}, Z={error_z:.2f}")
        print(f"Distance to target: {distance:.2f}m")
        print(f"Control inputs: ΔX={move_x:.2f}, ΔY={move_y:.2f}, ΔZ={move_z:.2f}")
        print("=" * 80)

        return move_x, move_y, move_z, 0

    def update_target(self):
        """更新目标点"""
        if self.reached_target and self.current_target_idx < len(self.target_positions) - 1:
            self.current_target_idx += 1
            self.reached_target = False

            # 打印目标更新信息
            target_x, target_y, target_z = self.target_positions[self.current_target_idx]
            print("\n" + "=" * 80)
            print(f"TARGET REACHED! MOVING TO TARGET {self.current_target_idx + 1}")
            print(f"New target: {target_x:.2f}, {target_y:.2f}, {target_z:.2f}")
            print("=" * 80)

            self.drone.log_event("MPC",
                                 f"Moving to target {self.current_target_idx}: {self.target_positions[self.current_target_idx]}")
            return True
        return False


# 演示函数
def run_demo():
    """运行完整的无人机MPC控制演示"""
    # 启动可视化
    drone.start_animation()

    # 创建MPC控制器
    mpc = MPCController(drone)

    # 记录开始时间
    start_time = time.time()

    # 执行演示序列
    print("Starting MPC control demonstration...")
    print("\nTARGET POINTS:")
    for i, (x, y, z) in enumerate(mpc.target_positions):
        print(f"T{i + 1}: ({x}, {y}, {z})")

    # 步骤1: 解锁无人机
    print("\n" + "=" * 80)
    print("STEP 1: ARMING DRONE")
    print("=" * 80)
    print(arm_drone())
    time.sleep(1)

    # 步骤2: 起飞
    print("\n" + "=" * 80)
    print("STEP 2: TAKEOFF")
    print("=" * 80)
    print(takeoff())
    time.sleep(2)

    # 步骤3: MPC控制循环
    print("\n" + "=" * 80)
    print("STEP 3: MPC CONTROL LOOP")
    print("=" * 80)
    demo_duration = 30  # 演示总时长（秒）
    control_interval = 0.5  # 控制间隔（秒）
    last_control_time = time.time()

    while time.time() - start_time < demo_duration:
        current_time = time.time()

        # 定期执行控制计算
        if current_time - last_control_time >= control_interval:
            last_control_time = current_time

            # 检查是否需要更新目标点
            if mpc.reached_target:
                if not mpc.update_target():
                    # 所有目标点已完成，准备降落
                    break

            # 计算控制输入
            control = mpc.calculate_control()
            if control:
                move_x, move_y, move_z, move_yaw = control
                print(move(move_x, move_y, move_z, move_yaw))

        # 更新图形
        plt.pause(0.01)

    # 步骤4: 返回起点并降落
    print("\n" + "=" * 80)
    print("STEP 4: RETURN TO HOME AND LAND")
    print("=" * 80)
    print(return_to_home())
    time.sleep(3)

    # 步骤5: 降落
    print("\n" + "=" * 80)
    print("STEP 5: LANDING")
    print("=" * 80)
    print(land())
    time.sleep(2)

    # 步骤6: 锁定无人机
    print("\n" + "=" * 80)
    print("STEP 6: DISARMING DRONE")
    print("=" * 80)
    print(disarm_drone())

    # 步骤7: 保存日志
    print("\n" + "=" * 80)
    print("STEP 7: SAVING LOGS")
    print("=" * 80)
    print(drone.save_logs_to_file("mpc_demo_log.json"))

    # 显示完成信息和统计
    print("\n" + "=" * 80)
    print("DEMONSTRATION COMPLETED!")
    print("=" * 80)
    print(f"Total flight time: {time.time() - start_time:.2f} seconds")
    print(f"Number of control calculations: {len(mpc.control_history)}")
    print(f"Final position: ({drone.x:.2f}, {drone.y:.2f}, {drone.z:.2f})")
    print("Visualization window will remain open. Close it to exit.")

    # 保持窗口打开
    plt.ioff()
    plt.show()


if __name__ == "__main__":
    run_demo()